# -*- coding: utf-8 -*-
from __future__ import absolute_import

"""
modis云掩膜产品处理
create_mask -of=mod35_cloudy_land.hdf -on=1 -off=0 -mask=”MOD35_L2.A1996213.1024.002.hdf,Cloud_Mask.1,1-2==00,AND,*,*,6-7==11” 
"""

from osgeo import gdal, osr, ogr
import numpy

import argparse

import sys
import os
import glob


def cloud_stats_month(shapefile, month_data_dir):
    """
    处理一个月的云掩膜数据给出报告
    :param month_data_dir:
    :return:
    """
    count = 0
    files = []
    files.extend(sorted(glob.glob("{}\\*Cloud_Mask_1.hdf".format(month_data_dir))))

    # new_files = []
    # for file in files:
    #     print(file)
    #     new_files.append(file.replace('\\', '/'))

    out_data_dir = "{}\\bin_cloud_mask".format(month_data_dir)
    if not os.path.isdir(out_data_dir):
        os.mkdir(out_data_dir)

    month_cloud_predicts = []
    datetime = []
    datetime.append("邮政编码")
    datetime.append("县级单位")
    for file in files:
        # print(file)
        datetime.append(file[42:49])
        count = count + 1
        out_bin_cloud = "{}\\{}.Binary_cloud.tif".format(out_data_dir, file[32:49])
        if not os.path.isfile(out_bin_cloud):
            bin_cloud_mask(file, out_bin_cloud)
        out_bin_cloud_cgcs = "{}_CGCS2000.tif".format(out_bin_cloud[:-4])
        if not os.path.isfile(out_bin_cloud_cgcs):
            cmd = "gdalwarp -s_srs EPSG:32649 -t_srs EPSG:4490 -r near -ot Byte -of GTiff {} {}"\
                .format(out_bin_cloud, out_bin_cloud_cgcs)
            os.system(cmd)
        stats_results = loop_feat_zonal_stats(shapefile, out_bin_cloud_cgcs) #返回一天的24个县市云量
        # print(stats_results)
        stats_results_arr = numpy.asarray(stats_results)
        # print(stats_results_arr[:, 2].astype(numpy.dtype(numpy.float32)))
        if count == 1:
            month_cloud_predicts.append(stats_results_arr[:, 0])
            month_cloud_predicts.append(stats_results_arr[:, 1])
        month_cloud_predicts.append(stats_results_arr[:, 2].astype(numpy.dtype(numpy.float32)))

    datetime_arr = numpy.asarray(datetime)
    month_cloud_predicts_arr = numpy.asarray(month_cloud_predicts).T
    output_statistics = numpy.vstack((datetime_arr, month_cloud_predicts_arr))
    numpy.savetxt("{}\\本月云量统计结果.csv".format(out_data_dir), output_statistics,
                  fmt='%s', delimiter=',')




def feat_zonal_stats(feat, input_zone_polygon, input_value_raster):
    """
    对一个feature进行栅格的区域统计
    :param feat:
    :param input_zone_polygon:这里传入的是矢量文件的路径
    :param input_value_raster:
    :return:
    """

    # Open data
    raster = gdal.Open(input_value_raster)
    shp = ogr.Open(input_zone_polygon)
    lyr = shp.GetLayer()
    #
    # # Get raster georeference info
    transform = raster.GetGeoTransform()
    xOrigin = transform[0]
    yOrigin = transform[3]
    pixelWidth = transform[1]
    pixelHeight = transform[5]
    #
    # Reproject vector geometry to same projection as raster
    sourceSR = lyr.GetSpatialRef()
    targetSR = osr.SpatialReference()
    targetSR.ImportFromWkt(raster.GetProjectionRef())
    coordTrans = osr.CoordinateTransformation(sourceSR, targetSR)

    geom = feat.GetGeometryRef()
    geom.Transform(coordTrans)

    city_pac = feat.GetField(1)  # PAC:469007
    city_name = feat.GetField(2)  #
    # Get extent of feat
    geom = feat.GetGeometryRef()
    if (geom.GetGeometryName() == 'MULTIPOLYGON'):
        count = 0
        pointsX = []; pointsY = []
        for polygon in geom:
            geomInner = geom.GetGeometryRef(count)
            ring = geomInner.GetGeometryRef(0)
            numpoints = ring.GetPointCount()
            for p in range(numpoints):
                    lon, lat, z = ring.GetPoint(p)
                    pointsX.append(lon)
                    pointsY.append(lat)
            count += 1
    elif (geom.GetGeometryName() == 'POLYGON'):
        ring = geom.GetGeometryRef(0)
        numpoints = ring.GetPointCount()
        pointsX = []; pointsY = []
        for p in range(numpoints):
                lon, lat, z = ring.GetPoint(p)
                pointsX.append(lon)
                pointsY.append(lat)

    else:
        sys.exit("ERROR: Geometry needs to be either Polygon or Multipolygon")

    #得到了输入polygon要素的最小外包矩形
    xmin = min(pointsX)
    xmax = max(pointsX)
    ymin = min(pointsY)
    ymax = max(pointsY)

    # Specify offset and rows and columns to read
    xoff = int((xmin - xOrigin)/pixelWidth)
    yoff = int((yOrigin - ymax)/pixelWidth)
    xcount = int((xmax - xmin)/pixelWidth)+1
    ycount = int((ymax - ymin)/pixelWidth)+1

    # Create memory target raster
    target_ds = gdal.GetDriverByName('MEM').Create('', xcount, ycount, 1, gdal.GDT_Byte)
    target_ds.SetGeoTransform((
        xmin, pixelWidth, 0,
        ymax, 0, pixelHeight,
    ))

    # Create for target raster the same projection as for the value raster
    raster_srs = osr.SpatialReference()
    raster_srs.ImportFromWkt(raster.GetProjectionRef())
    target_ds.SetProjection(raster_srs.ExportToWkt())

    # Rasterize zone polygon to raster
    gdal.RasterizeLayer(target_ds, [1], lyr, burn_values=[1]) #polygon外是0

    # Read raster as arrays
    #zone raster
    bandmask = target_ds.GetRasterBand(1)
    datamask = bandmask.ReadAsArray(0, 0, xcount, ycount).astype(numpy.float)

    #data raster to zone raster extent
    banddataraster = raster.GetRasterBand(1)
    dataraster = banddataraster.ReadAsArray(xoff, yoff, xcount, ycount).astype(numpy.float)


    # Mask zone of raster
    #掩膜设置填充值为-1， 方便后续用-1替换掩膜的--
    zoneraster = numpy.ma.masked_array(dataraster,  numpy.logical_not(datamask), fill_value=-1)

    # Calculate statistics of zonal raster
    unique, counts = numpy.unique(zoneraster, return_counts=True)
    unique_fill = unique.filled()
    count = 0
    count_cloud = 0
    count_nan = 0
    count_msk = 0
    for i in range(len(unique_fill)): #计数
        count += counts[i]
        if unique_fill[i] == 1:
            count_cloud = counts[i]
        if unique_fill[i] == -1:
            count_msk = counts[i]
        if unique_fill[i] == 3:
            count_nan = counts[i]
    cloud_cent = count_cloud / (count - count_msk)#计算百分比
    nan_cent = count_nan / (count - count_msk)

    # print(feat.GetField(1), feat.GetField(2), unique_fill, counts)

    # 这个条件判断是减小无数据对计算结果的干扰
    if nan_cent == 0: #不存在无数据，正常返回
        return feat.GetField(1), feat.GetField(2), cloud_cent
    elif nan_cent > 0.5: #无数据区域太大，不给出云量估计
        return feat.GetField(1), feat.GetField(2), -1
    else: #无数据区域不大，扣除无数据区给出云量估计
        return feat.GetField(1), feat.GetField(2), count_cloud / (count - count_nan - count_msk)

    # print(count, count_cloud, count_nan)
    # print(cloud_cent, nan_cent)
    # print(feat.GetField(1), feat.GetField(2), unique_fill, counts)

    # print(numpy.asarray((unique_fill, counts), dtype=numpy.dtype(numpy.int16)).T)

    #debug测试显示所用
    # plt.figure(1)
    # plt.subplot(211)
    # plt.title("{}".format(city_pac))
    # plt.imshow(datamask)
    # plt.subplot(212)
    # plt.imshow(dataraster)
    # plt.show()

    # return numpy.max(zoneraster)


def loop_feat_zonal_stats(input_zone_polygon, input_value_raster):
    """
    区域统计返回区域内最大值，并将值保存回原矢量layer中
    :param input_zone_polygon:
    :param input_value_raster:
    :return:返回更新后的矢量layer
    """
    shp = ogr.Open(input_zone_polygon, 0) #1 代表打开后可 update，默认是0，无法增加字段也无法写入数据
    lyr = shp.GetLayer()
    lyr.ResetReading()
    # print("feature count:", lyr.GetFeatureCount())
    count = 0
    feat = lyr.GetNextFeature() #PAC:469007 NAME:东方市

    # layerDefinition = lyr.GetLayerDefn()
    # for i in range(layerDefinition.GetFieldCount()):
    #     print(layerDefinition.GetFieldDefn(i).GetName())

    stats_results = []
    while feat:
        stats_result = feat_zonal_stats(feat, input_zone_polygon, input_value_raster)
        # print(stats_result)
        stats_results.append(stats_result)
        count = count + 1
        feat = lyr.GetNextFeature()
    return stats_results

def bin_cloud_mask(cloud_mask, off):
    """
    确定有云的为一类，其他的为一类，目的是给出类似“云量>10%”的建议
    :parameter cloud_mask 输入单个时间点的原始hdf文件
    :parameter off 输出文件
    """
    ds = gdal.Open(cloud_mask)
    bands_count = ds.RasterCount
    print(bands_count)
    prj = ds.GetProjection()
    srs = osr.SpatialReference(wkt=prj)
    band = ds.GetRasterBand(1)
    cloud_mask_arr = band.ReadAsArray()
    pans = ds.GetGeoTransform()
    bin_cloud_mask = numpy.zeros((cloud_mask_arr.shape[0], cloud_mask_arr.shape[1]))
    for i in range(cloud_mask_arr.shape[0]):
        for j in range(cloud_mask_arr.shape[1]):
            if cloud_mask_arr[i, j] != 0: #排除无数据区域
                if _bin(cloud_mask_arr[i, j])[-3: -1] == '00': #cloudy
                    bin_cloud_mask[i, j] = 1
                else:
                    # print(cloud_mask_arr[i, j])
                    bin_cloud_mask[i, j] = 2
            else:
                bin_cloud_mask[i, j] = 3  #no data
    array2raster(off, pans, bin_cloud_mask, srs)


def class_cloud_mask(cloud_mask, off):
    """
    云状况分类
    """
    ds = gdal.Open(cloud_mask)
    bands_count = ds.RasterCount
    print(bands_count)
    prj=ds.GetProjection()
    srs = osr.SpatialReference(wkt=prj)
    band = ds.GetRasterBand(1)
    cloud_mask_arr = band.ReadAsArray()
    pans = ds.GetGeoTransform()
    bin_cloud_mask = numpy.zeros((cloud_mask_arr.shape[0], cloud_mask_arr.shape[1]))
    for i in range(cloud_mask_arr.shape[0]):
        for j in range(cloud_mask_arr.shape[1]):
            if cloud_mask_arr[i, j] != 0: #排除无数据区域
                if _bin(cloud_mask_arr[i, j])[-3: -1] == '00': #cloudy
                    bin_cloud_mask[i, j] = 1
                elif _bin(cloud_mask_arr[i, j])[-3: -1] == '01': #uncertain clear
                    bin_cloud_mask[i, j] = 2
                elif _bin(cloud_mask_arr[i, j])[-3: -1] == '10': #probably clear
                    bin_cloud_mask[i, j] = 3
                elif _bin(cloud_mask_arr[i, j])[-3: -1] == '11': #confident clear
                    bin_cloud_mask[i, j] = 4
                else: #for nan
                    # print(cloud_mask_arr[i, j])
                    bin_cloud_mask[i, j] = 5
            else:
                bin_cloud_mask[i, j] = -999
    array2raster(off, pans, bin_cloud_mask, srs)



def _bin(dig):
    """
    给python内置的bin函数补位
    :param dig:
    :return:
    """
    return bin(dig)[2:].zfill(8)


def array2raster(newRasterfn, panTransform, array, srs):
    """
    将数组保存为UTM投影的遥感影像
    :param newRasterfn:
    :param panTransform: imggt
    :param array:
    :return:
    """
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((panTransform[0], panTransform[1], panTransform[2], panTransform[3],
                               panTransform[4], panTransform[5]))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = srs
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()


def main():
    parser = argparse.ArgumentParser(description="统计一个月的云情况，尺度为全省24个县级行政单位")
    parser.add_argument("shpfile", type=str, help="当地polygon或者multipolygon shapefile文件")
    parser.add_argument("data_dir", type=str, help="MODIS云掩膜产品按月存储的数据目录")
    # # parser.add_argument("uav", type=str, help="path to uav image")
    # # parser.add_argument("-w", type=int, help="search circul while compute diff map")
    # # parser.add_argument("-yuzhi_method", type=str, choices=['rosin', 'polygon', 'mouse'], help="select a yuzhi method for diffmap binary")
    # # parser.add_argument("-min_area", type=int, help="min pixels of a PCC patch")
    args = parser.parse_args()
    # sParams = SParams._make([args.sat, args.uav, args.w, args.yuzhi_method, args.min_area])
    cloud_stats_month(args.shpfile, args.data_dir)
    # cloud_stats_month("F:\\data\\modis_cloud_mask\\shapefile\\hainan_24_county.shp",
    #                   "F:\\data\\modis_cloud_mask\\201711")


if __name__ == '__main__':
    main()